# -*- coding: utf-8 -*-
# @Author: Cursor
# @Date: 2025-01-08
# @Last Modified by: Tim Liu
# @Last Modified time: 2025-01-08

# @Description: Ontology DAL for RAG

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import joinedload
from sqlalchemy import select 
from core.crud import DalBase
from crewplus.apps.rag import models, schemas
from crewplus.apps.rag.models.ontology_schema_association import ontology_schema_association

class OntologyDal(DalBase):

    def __init__(self, db: AsyncSession):
        super(OntologyDal, self).__init__()
        self.db = db
        self.model = models.OntologyDB
        self.schema = schemas.OntologySimpleOut
        
    async def add_oschema_to_ontology(self, ontology_id: int, oschema_id: int) -> None:
        """
        Add an OSchema to the ontology.
        """
        ontology = await self.get_data(ontology_id)
        oschema = await self.db.get(models.OSchemaDB, oschema_id)
        
        if not ontology or not oschema:
            raise ValueError("Invalid ontology or schema ID")

        # Check if the association already exists
        association_exists = await self.db.execute(
            select(ontology_schema_association).where(
                ontology_schema_association.c.ontology_id == ontology_id,
                ontology_schema_association.c.schema_id == oschema_id
            )
        )
        if association_exists.scalar():
            raise ValueError("Association already exists")

        # Add the association
        await self.db.execute(
            ontology_schema_association.insert().values(
                ontology_id=ontology_id,
                schema_id=oschema_id
            )
        )
        await self.db.commit()
        
    async def add_multiple_schemas_to_ontology(self, ontology_id: int, schema_ids: list[int]) -> None:
        """
        Add multiple schemas to the ontology.
        """
        ontology = await self.get_data(ontology_id)
        if not ontology:
            raise ValueError("Ontology not found")

        schemas = await self.db.execute(select(models.OSchemaDB).where(models.OSchemaDB.id.in_(schema_ids)))
        schemas = schemas.scalars().all()

        if not schemas:
            raise ValueError("Schemas not found")

        for schema in schemas:
            # Check if the association already exists
            association_exists = await self.db.execute(
                select(ontology_schema_association).where(
                    ontology_schema_association.c.ontology_id == ontology_id,
                    ontology_schema_association.c.schema_id == schema.id
                )
            )
            if not association_exists.scalar():
                # Add the association
                await self.db.execute(
                    ontology_schema_association.insert().values(
                        ontology_id=ontology_id,
                        schema_id=schema.id
                    )
                )

        await self.db.commit()
        
    async def get_ontology_with_schemas(self, ontology_id: int) -> schemas.OntologyFullOut:
        """
        Get ontology information, including associated schemas.
        """
        ontology = await self.db.execute(
            select(models.OntologyDB)
            .options(joinedload(models.OntologyDB.schemas))
            .where(models.OntologyDB.id == ontology_id)
        )
        ontology = ontology.scalar_one_or_none()

        if not ontology:
            raise ValueError("Ontology not found")

        return self.schema.model_validate(ontology).model_dump()